#!/usr/bin/env python3
'''
    单独测试脚本，用于测试网络监听的准确度和平均耗时
    采用连接100次，计算平均耗时和成功率
    测试结果：
    成功率：96/100
    avg:：39s
'''
from scapy.all import Dot11, sniff, conf, raw
import os
import sys
import time
import subprocess
import threading
import socket
from datetime import datetime
import json

# 常量定义
INTERFACE = "wlan0mon"
BROADCAST_MAC = "ff:ff:ff:ff:ff:ff"
HEADER_LEN = 361
ASCII_START = 127
ASCII_END = 128

# 全局状态变量
is_active = False
frame_data = []
data_dict = {}
total_len = 0
total_len_list = []
total_len_count = 0
current_index = -1
target_data = []
sniffing_thread = None
sniffing_event = threading.Event()
exit_event = threading.Event()  # 新增退出信号
has_check_repeat_flag = False

CHECK_NUMS=2
check_num = 0
vaild_num = 0 
vaild_wifi='HUAWEI-7EGBK8_HiLink:1231231230'



def check_root():
    if os.geteuid() != 0:
        print("❌ 错误：需要root权限！")
        sys.exit(1)

def ascii_char(code: int) -> str:
    if 32 <= code <= 126:
        return chr(code)
    return f"(非打印:{code})"

def reset():
    global is_active, frame_data, has_check_repeat_flag
    is_active = False
    frame_data.clear()
    has_check_repeat_flag = False

def parse_wifi_credentials(credentials_str):
    try:
        ssid, password = credentials_str.split(':', maxsplit=1)
        return ssid.strip(), password.strip()
    except ValueError:
        raise ValueError("无效的凭证格式，应为 'SSID:password'")

def switch_to_managed(monitor_iface="wlan0mon", managed_iface="wlan0"):
    global INTERFACE
    print(f"📡 停止监控模式接口: {monitor_iface}")
    subprocess.run(
        ["airmon-ng", "stop", monitor_iface],
        check=True,
        capture_output=True,
        text=True
    )
    print(f"🔄 设置接口 {managed_iface} 为 managed 模式")
    subprocess.run(
        ["ip", "link", "set", managed_iface, "down"],
        check=True,
        capture_output=True,
        text=True
    )
    subprocess.run(
        ["iw", managed_iface, "set", "type", "managed"],
        check=True,
        capture_output=True,
        text=True
    )
    subprocess.run(
        ["ip", "link", "set", managed_iface, "up"],
        check=True,
        capture_output=True,
        text=True
    )
    print("🔄 启动 NetworkManager")
    subprocess.run(
        ["systemctl", "start", "NetworkManager"],
        check=True,
        capture_output=True,
        text=True
    )
    time.sleep(5)
    INTERFACE = managed_iface
    print(f"✅ 已切换到 managed 模式: {managed_iface}")
    

def switch_to_monitor(managed_iface="wlan0", monitor_iface="wlan0mon", timeout=10):
    """切换为监控模式（带启动状态检测）"""
    global INTERFACE
    
    print("🔄 停止 NetworkManager")
    subprocess.run(
        ["airmon-ng", "check", "kill"],
        check=True,
        capture_output=True,
        text=True
    )
    time.sleep(3)
    print(f"📡 尝试启动监控模式: {monitor_iface}")
    try:
        # 2. 启动监控模式（指定频道6）
        subprocess.run(
            ["airmon-ng", "start", managed_iface, "6"],
            check=True,
            capture_output=True,
            text=True
        )
        
        # 3. 循环检测接口状态（带超时）
        start_time = time.time()
        while time.time() - start_time < timeout:
            # 检测接口是否存在
            if not os.path.exists(f"/sys/class/net/{monitor_iface}"):
                time.sleep(1)
                continue
            print(f"{monitor_iface}存在！")
            # 检测监控模式是否激活
            result = subprocess.run(["iwconfig", monitor_iface], 
                                  capture_output=True, text=True)
            if "Mode:Monitor" in result.stdout:
                INTERFACE = monitor_iface
                print(f"✅ 监控模式启动成功 | 接口: {monitor_iface} | 频道: 6")
                return True
            time.sleep(1)

        print(f"❌ 超时：{monitor_iface} 未在 {timeout} 秒内进入监控模式")
        return False
        
    except subprocess.CalledProcessError as e:
        print(f"❌ 启动失败: {e.stderr.strip()}")
        return False


def broadcast_ack(port=12345, ack_message="ACK"):
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        sock.settimeout(0.5)
        message = ack_message.encode('utf-8')
        for i in range(10):
            sock.sendto(message, ('255.255.255.255', port))
            print(f"✅ 已广播ACK到255.255.255.255:{port}")
    except Exception as e:
        print(f"❌ 广播失败: {str(e)}")
    finally:
        sock.close()

def start_sniffing():
    global INTERFACE
    while not exit_event.is_set():  # 检查退出信号
        if sniffing_event.is_set():
            try:
                print(f"🚀 开始嗅探接口: {INTERFACE}")
                sniff(iface=INTERFACE, prn=handle_packet, store=0, monitor=True, timeout=1)
            except Exception as e:
                print(f"❌ 嗅探错误: {e}")
        else:
            print("🔄 嗅探暂停，等待重启指令...")
            time.sleep(1)
    print("🛑 嗅探线程退出")


def handle_wifi_connection(wifi_string, sniffing_event, exit_event):
    """
    处理WiFi连接流程（包含模式切换、连接和错误恢复）
    
    参数:
        target_data: 收集到的凭证数据列表
        sniffing_event: 线程事件控制嗅探状态
        exit_event: 线程事件控制程序退出
    """

    try:
        # 2. 解析凭证
        ssid, password = parse_wifi_credentials(wifi_string)
        print(f"🔑 尝试连接: SSID='{ssid}'")

        # 3. 执行连接流程
        sniffing_event.clear()  # 暂停嗅探
        if  _establish_wifi_connection(ssid, password):
            broadcast_ack()
            print("✅ 连接成功，程序退出")
            exit_event.set()
        else:
            raise RuntimeError("连接流程失败")

    except subprocess.CalledProcessError as e:
        print(f"❌ 连接失败: {e.stderr.strip()}")
        _recover_to_monitor(sniffing_event, exit_event)

    except ValueError as e:
        print(f"❌ 凭证无效: {e}")
        sniffing_event.set()  # 恢复嗅探

    except Exception as e:
        print(f"‼️ 未知错误: {e}")
        _recover_to_monitor(sniffing_event, exit_event)

    print("=" * 50)

def _establish_wifi_connection(ssid, password, monitor_iface="wlan0mon", managed_iface="wlan0"):
    """执行实际的WiFi连接流程"""
    global INTERFACE

    # 1. 切换到托管模式
    switch_to_managed(monitor_iface, managed_iface)

    # 2. 执行连接命令
    result = subprocess.run(
        ["nmcli", "device", "wifi", "connect", ssid, "password", password, "ifname", managed_iface],
        check=True,
        capture_output=True,
        text=True
    )
    
    # 3. 验证连接结果
    if "successfully" not in result.stdout.lower():
        raise RuntimeError("连接响应异常")
    
    INTERFACE = managed_iface
    return True

def _recover_to_monitor(sniffing_event, exit_event):
    """恢复监控模式并处理异常状态"""
    if not switch_to_monitor():
        print("❌ 无法恢复监控模式!")
        exit_event.set()
    else:
        sniffing_event.set()  # 恢复嗅探


def find_duplicate_string(target_data):
    global has_check_repeat_flag
    """
    检查target_data列表中是否存在重复字符串
    
    参数:
        target_data: 字符串列表
        
    返回:
        str: 返回第一个被发现的重复字符串
        None: 如果没有重复则返回None
    """
    seen = set()
    for item in target_data:
        if item in seen:
            return item  # 返回第一个重复的字符串
        seen.add(item)
    has_check_repeat_flag = True
    return None  # 没有重复


def handle_packet(pkt):
    global is_active, frame_data, data_dict, total_len, total_len_list, total_len_count
    global current_index, target_data, sniffing_event, exit_event, has_check_repeat_flag
    global CHECK_NUMS, check_num, vaild_num, vaild_wifi
    if not (pkt.haslayer(Dot11) and pkt.type == 2 and pkt.addr1 == BROADCAST_MAC):
        return
    
    frame_len = len(raw(pkt))
    if frame_len < HEADER_LEN:
        return
    
    ascii_code = frame_len - HEADER_LEN
    timestamp = time.strftime("%H:%M:%S", time.localtime())
    print(f"\n🌐 [{timestamp}] 捕获帧 (长度: {ascii_code} 字节)")
    
    frame_data.append(ascii_code)
    frame_data_len = len(frame_data)
    if ascii_code == ASCII_START:
        if frame_data_len > 1:
            # 开始帧不在首位，帧读取错误，清空数据，以当前开始帧重新开始
            reset()
            frame_data.append(ascii_code)
        is_active = True
        return
    
    if ascii_code == ASCII_END:
        print("  🛑 检测到结束帧")
        print(f"  📋 帧数据: {frame_data}")
        print(f"  📝 解码结果: {data_dict}")
        print(f"  🔢 总长度计数: {total_len_count}")
        print(f"  📊 总长度候选: {total_len_list}")
        print(f"  📏 最终总长度: {total_len}")
        print(f"  🎉 解码结果集合: {target_data}")
        print('-' * 50)
        reset()
        return
    
    if not is_active:
        return
    
    if len(frame_data) > 5:
        print("  ⚠️ 帧数量过多，重置")
        reset()
        return
    
    if len(frame_data) == 2:
        temp_len = ascii_code
        print(f"  📏 当前总长度: {temp_len}")
        if total_len_count < 5:
            total_len_list.append(temp_len)
            total_len_count += 1
        elif total_len_count >= 5 and total_len == 0:
            total_len = max(set(total_len_list), key=total_len_list.count)
            total_len_list.clear()
            print(f"  ✅ 确认总长度: {total_len}")
    
    elif len(frame_data) == 3:
        current_index = ascii_code
    
    elif len(frame_data) == 4:
        char = ascii_char(ascii_code)
        print(f"  🔤 解码字符: {char}")
        if current_index >= 0 and len(char) == 1:
            data_dict[current_index] = char
            print(f"data_dict: {data_dict}")
        else:
            print("  ❌ 解码失败，丢弃帧")
    
    if total_len > 0:
        if len(data_dict) == total_len:
            result = "".join(data_dict[k] for k in sorted(data_dict.keys()))
            print(f"  🎉 解码完成: {result}")
            target_data.append(result)
            has_check_repeat_flag = False
            data_dict.clear()
            reset()
        elif len(data_dict) > total_len:
            result = "".join(data_dict[k] for k in sorted(data_dict.keys())[:total_len])
            print(f"  🎉 存在多余数据，进行截取，解码完成: {result}")
            target_data.append(result)
            has_check_repeat_flag = False
            data_dict.clear()
            reset()

    if len(target_data) > 1 and not has_check_repeat_flag:
        # 如果存在两个一样的字符串，则直接发送这个字符串
        wifi_string = find_duplicate_string(target_data)
        if wifi_string:
            print("快捷算法：有两个相同的字符串直接发送！")
            target_data.clear()
            # handle_wifi_connection(wifi_string, sniffing_event, exit_event)
            if vaild_wifi == wifi_string:
                vaild_num += 1
            check_num += 1
            if check_num >= CHECK_NUMS:
                print(f"vaild times:{vaild_num}/{check_num}")
                exit_event.set()


    if len(target_data) > 5:
        # 1. 提取出现频率最高的凭证
        wifi_string = max(set(target_data), key=target_data.count)
        print('*' * 50)
        print(f"  🎉 最佳候选凭证: {wifi_string}")
        print('*' * 50)
        target_data.clear()

        if vaild_wifi == wifi_string:
            vaild_num += 1
        check_num += 1
        if check_num >= CHECK_NUMS:
            print(f"vaild times:{vaild_num}/{check_num}")
            exit_event.set()
        # handle_wifi_connection(wifi_string, sniffing_event, exit_event)

    print("=" * 50)


def set_channel(iface, channel):
    """设置无线接口到指定信道"""
    try:
        print(f"🔧 设置接口 {iface} 到信道 {channel}")
        subprocess.run(
            ["iw", iface, "set", "channel", str(channel)],
            check=True,
            capture_output=True,
            text=True
        )
        time.sleep(1)  # 等待信道切换稳定
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ 设置信道 {channel} 失败: {e.stderr.strip()}")
        return False

def scan_channels(iface="wlan0mon", channels=[1, 6, 11], duration=1):
    """扫描指定信道，统计每个信道的包数量，并将包最多的信道设置为当前信道"""
    packet_counts = {ch: 0 for ch in channels}
    
    def packet_counter(pkt):
        """回调函数，统计数据包数量"""
        if pkt.haslayer(Dot11):
            packet_counts[current_channel] += 1
    
    for channel in channels:
        if not set_channel(iface, channel):
            print(f"❌ 跳过信道 {channel} 的扫描")
            continue
        
        print(f"📡 开始扫描信道 {channel}，持续 {duration} 秒")
        try:
            conf.use_pcap = True  # 确保使用 libpcap
            current_channel = channel
            packet_counts[channel] = 0
            sniff(iface=iface, prn=packet_counter, store=0, monitor=True, timeout=duration)
            print(f"✅ 信道 {channel} 扫描完成，接收到 {packet_counts[channel]} 个数据包")
        except Exception as e:
            print(f"❌ 信道 {channel} 扫描错误: {e}")
    
    # 选择包数量最多的信道
    if packet_counts:
        best_channel = max(packet_counts, key=packet_counts.get)
        max_packets = packet_counts[best_channel]
        print(f"\n🔍 最佳信道: {best_channel} (接收到 {max_packets} 个数据包)")
        
        # 设置为最佳信道
        if set_channel(iface, best_channel):
            print(f"✅ 接口 {iface} 已设置为信道 {best_channel}")
        else:
            print(f"❌ 无法将接口 {iface} 设置为信道 {best_channel}，保持当前信道")
    else:
        print("\n❌ 所有信道扫描失败，无法选择最佳信道")
        best_channel = None
    
    return packet_counts, best_channel

def main():
    check_root()
    print(f"🎯 目标MAC: {BROADCAST_MAC}")
    print("按 Ctrl+C 停止")
    
    global sniffing_thread, sniffing_event, exit_event, check_num, vaild_num
    try:
        # use_pcap非常重要，启用使用libcap比起原始套接字更兼容和稳定！如果不配置，会导致无法重启嗅探！
        conf.use_pcap = True
        # 扫描信道，确定发信源信道
        scan_channels()
        start = time.time()
        # 正式嗅探捕获数据
        sniffing_event.set()
        sniffing_thread = threading.Thread(target=start_sniffing, daemon=True)
        sniffing_thread.start()

        while not exit_event.is_set():  # 监控退出信号
            time.sleep(1)

        end = time.time()
        elapsed = end - start
        minutes = int(elapsed // 60)
        seconds = int(elapsed % 60)
        print(f"vaild times:{vaild_num}/{check_num}")
        print(f"cost: {minutes}分钟{seconds}秒")  # 输出示例：0分钟3秒
        print(f"avg:{elapsed/check_num}秒")
        print("🛑 主线程退出")
    except KeyboardInterrupt:
        print("\n🛑 用户停止监听")
        exit_event.set()

if __name__ == "__main__":
    main()